import angr
import sys

import angr.knowledge_base

sys.path.append('src')

from tasks.lsbvd.feature.const_value import buffer_overflow_funcs,buffer_overflow_sink_funcs,buffer_overflow_source_funcs


class CallRelationShipFeatures:
    def __init__(self, cfg_fast:angr.analyses.cfg.cfg_fast.CFGFast):
        self.cfg= cfg_fast

    def get_func_calls(self, func_addrs):
        """
        获取函数调用点
        :param func_addrs: 函数地址列表
        :return: 函数调用点列表
        """
        return get_func_calls(self.cfg, func_addrs)

class SinkOrientSearcher:
    def __init__(self,proj:angr.Project, cfg_fast:angr.analyses.cfg.cfg_fast.CFGFast,name2addr:dict=None):
        self.cfg = cfg_fast
        self.proj = proj
        self.name2addr=name2addr
        self.d0_sink=self._get_d0_list(buffer_overflow_sink_funcs())
        self.d0_source=self._get_d0_list(buffer_overflow_source_funcs())
        
        
        
        self.depth_caches_source={0:self.d0_source}
        self.depth_caches_sink={0:self.d0_sink}

    def source_sink_pair(self):
        source_paths=[]
        for source in self.d0_source:
            source_paths+=self.call_path(source.addr)
            
        sink_paths=[]
        for sink in self.d0_sink:
            sink_paths+=self.call_path(sink.addr)
            
        #pairs=[]
        for source_path in source_paths:
            for sink_path in sink_paths:
                common_nodes=self.common_nodes(source_path,sink_path)
                if len(common_nodes)>0:
                    yield source_path,sink_path,common_nodes
    
    def generate_source_paths(self):
        """
        生成源函数调用路径
        :return: 源函数调用路径列表
        """
        for source in self.d0_source:
            paths = self.call_path(source.addr)
            for path in paths:
                yield path
    
    def generate_sink_paths(self):
        """
        生成汇函数调用路径
        :return: 汇函数调用路径列表
        """
        for sink in self.d0_sink:
            paths = self.call_path(sink.addr)
            for path in paths:
                yield path
    
    
    def _get_d0_list(self,func_names):
        candidate_names=func_names
        
        ret=[]
        
        for name in candidate_names:
            # 获取函数地址
            func = self.proj.kb.functions.function(addr=self.name2addr[name])
            if func is not None:
                # 获取函数的地址
                #addr = func.addr
                # 获取函数的调用点
                # site_number = self.call_graph.in_degree(addr)
                # if site_number > 0:
                ret.append(func)
                    
        return ret

    def depth_n_sink(self,n:int):
        return self.depth_n(n,self.depth_caches_sink)
    
    def depth_n_source(self,n:int):
        return self.depth_n(n,self.depth_caches_source)

    def depth_n(self,n:int,caches):
        """
        获取深度为n的函数
        :param n: 深度
        :return: 函数列表
        """
        if n in caches:
            return caches[n]
        
        
        previous_depth=self.depth_n(n-1,caches)
        
        
        ret=set()
        
        for func in previous_depth:
            # 获取函数的调用点
            try:
                predecessors = self.call_graph.predecessors(func.addr)
                for predecessor in predecessors:
                    if predecessor not in previous_depth:
                        ret.add(self.proj.kb.functions.function(addr=predecessor))
            except:
                # 如果函数不存在,则跳过
                continue
                
        self.depth_caches_sink[n]=list(ret)
        
        return ret

    def call_path(self, func_addr:int,max_number=1000000):
        # 以call_site为起点,深度遍历call_graph
        
        all_paths=[]
        
        def get_path(node,path:list,visited:set):
            visited.add(node)
            path.append(node)
            
            if len(all_paths) >= max_number:
                return
            
            # 获取函数的调用点
            
            if node not in self.call_graph:
                all_paths.append(path.copy())
                path.pop()
                visited.remove(node)
                return
            
            preds = list(self.call_graph.predecessors(node))
            if len(preds) == 0:
                all_paths.append(path.copy())
            else:
                for pred in preds:
                    if pred not in visited:
                        get_path(pred, path, visited)
            
            path.pop()
            visited.remove(node)
        
        get_path(func_addr,[],set())
        
        
        return all_paths
    
    def common_nodes(self, path1, path2):

        common_nodes = set(path1) & set(path2)
        
        return list(common_nodes)
    
    @property
    def func_number(self):
        """
        获取函数数量
        :return: 函数数量
        """
        return len(self.cfg.kb.functions)


    @property
    def call_graph(self):
        """
        获取函数调用图
        :return: 函数调用图
        """
        return self.cfg.kb.functions.callgraph

    def all_funcs(self):
        """
        获取所有函数
        :return: 函数列表
        """
        return self.cfg.kb.functions.values()


if __name__ == '__main__':
    import time
    import json
    from tasks.lsbvd.feature.static_cfg import BigBinaryCfgBuilder
    
    region_split_path='caches_data/amd.txtbin_partitions.json'

    # region_split by user define region
    region_json=json.load(open(region_split_path, 'r'))['partitions']
    
    bbinary_builder= BigBinaryCfgBuilder(
        binary_path='caches_data/ine.txtbin',
        caches_dir='caches_data/nokia_cfg',
        symbols_json_path='caches_data/symbols_export.json'
    )
    # proj=bbinary_builder.load_project()
    # funcs=proj.kb.functions.values()
    # for e in funcs:
    #     func=e
    #     if func.name.startswith('sub'):
    #         continue
    #     print(f'func: {func.name} {hex(func.addr)}')
    
    cfgs=bbinary_builder.load_partial_cfg(regions=region_json)
    
    cfg=None
    for e in cfgs:
        cfg=e
        break
    
    # print(f'Loaded {len(cfgs)} cfgs')
    
        
    print("process cfg.")
    
    proj=bbinary_builder.proj
    
    
    searcher=SinkOrientSearcher(bbinary_builder.proj,cfg,bbinary_builder.name2iddr)
        
    
    source_counter=0
    sink_counter=0
    
    with open('caches_data/nokia_path/source_path.txt','w') as f:
        for path in searcher.generate_source_paths():
            source_counter+=1
            if source_counter % 1000 == 0:
                print(f'Processed {source_counter} source paths')
            f.write(','.join([str(addr) for addr in path])+'\n')
            
    with open('caches_data/nokia_path/sink_path.txt','w') as f:
        for path in searcher.generate_sink_paths():
            sink_counter+=1
            if sink_counter % 1000 == 0:
                print(f'Processed {sink_counter} sink paths')
            f.write(' '.join([str(addr) for addr in path])+'\n')    
    
        # funcs=searcher.all_funcs()
        
        # for i,e in enumerate(funcs):
        #     func=e
        #     print(f'func {i}: {func.name} {hex(func.addr)}')
        #     if i>100:
        #         break
        
        # d0_sink=searcher.depth_n_sink(0)
        # d0_source=searcher.depth_n_source(0)
        
        # for func in d0_sink:
        #     print(f'sink func: {func.name} {hex(func.addr)}')
        # for func in d0_source:
        #     print(f'source func: {func.name} {hex(func.addr)}')
        
        # print('------------------')
            
        # d1_sink=searcher.depth_n_sink(1)
        # d1_source=searcher.depth_n_source(1)
        
        # for func in d1_sink:
        #     if func:
        #         print(f'sink func: {func.name} {hex(func.addr)}')
        # for func in d1_source:
        #     if func:
        #         print(f'source func: {func.name} {hex(func.addr)}')
        
        
        # pairs=searcher.source_sink_pair()
        # for pair in pairs:
        #     source_path,sink_path,common_nodes=pair
            
        #     source_names=[]
        #     sink_names=[]
        #     common_nodes_names=[]
            
        #     for common_node in common_nodes:
        #         common_nodes_names.append(cfg.kb.functions.function(addr=common_node).name)
            
        #     for source in source_path:
        #         source_names.append(cfg.kb.functions.function(addr=source).name)
                
        #     for sink in sink_path:
        #         sink_names.append(cfg.kb.functions.function(addr=sink).name)
                
            
        #     print('source:',source_names)
        #     print('sink:',sink_names)
        #     print('common:',common_nodes_names)
        #     print('------------------')
    
    
    # 获取加载的内存范围
    # import struct
    # mem = proj.loader.memory
    # # 假设 .data 段前 4 字节是一个 32 位整数
    # first_word = mem.load(proj.loader.main_object.min_addr, 4)
    # value = struct.unpack("<I", first_word)[0]  # 小端序解析
    # print(f"First 32-bit value: {hex(value)}")
    # cfg=proj.analyses.CFGFast()
    
    # print('cfg time:',time.time()-start)
    
    # searcher=SinkOrientSearcher(proj,cfg)
    
    # pairs=searcher.source_sink_pair()
    # for pair in pairs:
    #     source_path,sink_path,common_nodes=pair
        
    #     source_names=[]
    #     sink_names=[]
    #     common_nodes_names=[]
        
    #     for common_node in common_nodes:
    #         common_nodes_names.append(cfg.kb.functions.function(addr=common_node).name)
        
    #     for source in source_path:
    #         source_names.append(cfg.kb.functions.function(addr=source).name)
            
    #     for sink in sink_path:
    #         sink_names.append(cfg.kb.functions.function(addr=sink).name)
            
        
    #     print('source:',source_names)
    #     print('sink:',sink_names)
    #     print('common:',common_nodes_names)
    #     print('------------------')
# def get_func_calls(cfg,func_addrs):
#     """
#     获取函数调用点
#     :param cfg: CFG对象
#     :param name: 函数名称
#     :return: 函数调用点列表
#     """

#     #functions=[proj.kb.functions.function(addr=addr) for addr in func_addrs]

#     if not isinstance(func_addrs, list):
#         func_addrs = [func_addrs]

#     func_calls = []
#     for addr, func in cfg.kb.functions.items():
#         # 获取函数的调用图
#         for call_site in func.get_call_sites():
#             # 获取被调用的函数
#             target = func.get_call_target(call_site)
#             if target in func_addrs:
#                 func_calls.append(call_site)

#     return func_calls